"""
This part of code is the Q learning brain, which is a brain of the agent.
All decisions are made in here.

Reference: https://morvanzhou.github.io/tutorials/
"""

import numpy as np
import pandas as pd


class QLearningTable:
    def __init__(self, actions, learning_rate=0.01, reward_decay=0.9, e_greedy=0.9):  # 初始化
        self.actions = actions  # a list
        self.lr = learning_rate # 学习率
        self.gamma = reward_decay # 衰减率/折扣因子
        self.epsilon = e_greedy # 贪婪度
        self.q_table = pd.DataFrame(columns=self.actions, dtype=np.float64) # 空的Q表

    def choose_action(self, observation):  # 选择动作
        self.check_state_exist(observation)  # 检测本state是否在q表中存在（不存在要添加）
        # action selection
        if np.random.uniform() < self.epsilon:  # 选择Q值最高的动作
            # choose best action
            state_action = self.q_table.loc[observation, :]  # 找到本state的所有action值
            # some actions may have the same value, randomly choose on in these actions
            action = np.random.choice(state_action[state_action == np.max(state_action)].index) 
            # 有多个相同的Q值时随机选一个
        else:
            # choose random action 
            action = np.random.choice(self.actions)  # 有一定概率随机探索选取一个动作
        return action

    def learn(self, s, a, r, s_): # 学习更新Q表
        self.check_state_exist(s_)  # 检测q表中是否存在s_
        q_predict = self.q_table.loc[s, a] 
        if s_ != 'terminal':
            q_target = r + self.gamma * self.q_table.loc[s_, :].max()  # next state is not terminal
        else:
            q_target = r  # next state is terminal
        self.q_table.loc[s, a] += self.lr * (q_target - q_predict)  # update

    def check_state_exist(self, state):  # 检测本state是否在q表中存在
        '''
        #append方法废止
        if state not in self.q_table.index:
            # append new state to q table
            self.q_table = self.q_table._append( # 添加新的一行
                pd.Series(
                    [0]*len(self.actions),
                    index=self.q_table.columns, 
                    name=state, 
                )
            )    def check_state_exist(self, state):  # 检测本state是否在q表中存在
        '''
        if state not in self.q_table.index:
            # append new state to q table
            new_state = pd.Series(
                [0]*len(self.actions),
                index=self.q_table.columns,
                name=state,
            )
            self.q_table = pd.concat([self.q_table, new_state.to_frame().T])  